import json
import pickle
import os
import torch

import numpy as np
import h5py
import utils
from torch.utils.data import Dataset

class Dictionary(object):
    def __init__(self, word2idx=None, idx2word=None):
        if word2idx is None:
            word2idx = {}
        if idx2word is None:
            idx2word = []
        self.word2idx = word2idx
        self.idx2word = idx2word

    @property
    def ntoken(self):
        return len(self.word2idx)

    @property
    def padding_idx(self):
        return len(self.word2idx)

    def tokenize(self, sentence, add_word):
        '''
        将句子转化为字典id。
        Args:
            sentence(str): 待输入句子
            add_word(bool): 句子中是否含有新单词。若是，则将单词加入到索引中。

        Returns:
            list[int]: 句子的字典编码

        '''
        sentence = sentence.lower()
        sentence = sentence.replace(',', '').replace('?', '').replace("'s", " 's")
        words = sentence.split()
        tokens = []

        if add_word:
            for w in words:
                tokens.append(self.add_word(w))
        else:
            for w in words:
                tokens.append(self.word2idx[w])
        return tokens

    def dump_to_file(self, path):
        '''
        将字典写入到文件。
        Args:
            path(str): 保存的路径

        Returns:

        '''
        pickle.dump([self.word2idx, self.idx2word], open(path, 'wb'))
        print('dictionary dumped to {}'.format(path))

    @classmethod
    def load_from_file(cls, path):
        '''
        从文件中加载字典
        Args:
            path(str): 文件路径

        Returns:
            Dictionary: 构造一个词典
        '''
        print('Loading dictionary from %s' % path)
        word2idx, idx2word = pickle.load(open(path, 'rb'))
        d = cls(word2idx, idx2word)
        return d

    def __len__(self):
        return len(self.idx2word)

    def add_word(self, word):
        '''
        将新单词加入到索引中
        Args:
            word(str): 要加入的单词

        Returns:
            int: 该单词的下标。
        '''
        if word not in self.word2idx:
            self.idx2word.append(word)
            self.word2idx[word] = len(self.idx2word) - 1
        return self.word2idx[word]


def _load_dataset(dataroot, name, img_id2val):
    '''
    Load entries

    Args:
        dataroot(str): data root path
        name(str): train or val
        img_id2val(dict): image id to id(img_id为coco中的图片标号，不连续；id为features中对应的标号)

    Returns:
        List: entries
    '''
    question_path = os.path.join(
        dataroot, 'v2_OpenEnded_mscoco_{}2014_questions.json'.format(name)
    )
    questions = sorted(json.load(
        open(question_path, 'r')['questions'],
        key=lambda x: x['question_id']
    ))
    answer_path = os.path.join(
        dataroot, 'cache', '{}_target.pkl'.format(name)
    )
    answers = pickle.load(open(answer_path, 'rb'))
    answers = sorted(answers, key=lambda x:x['question_id'])

    utils.assert_eq(len(questions), len(answers))
    entries = []
    for question, answer in zip(questions, answers):
        utils.assert_eq(question['qustion_id'], answer['question_id'])
        utils.assert_eq(question['image_id'], answer['image_id'])
        img_id = question['image_id']
        entries.append(_create_entry(img_id2val[img_id], question, answer))

def _create_entry(img, question, answer):
    answer.pop('image_id')
    answer.pop('question_id')
    entry = {
        'question_id' : question['question_id'],
        'image_id'    : question['image_id'],
        'image'       : img,
        'question'    : question['question'],
        'answer'      : answer
    }
    return answer



class VQAFeatureDataset(Dataset):

    def __init__(self, name, dictionary, dataroot='data'):
        '''
        Provide VQA Features. Including visual features and question features.

        Args:
            name(str): train or val
            dictionary(Dictionary): A dictionary object.
            dataroot(str): The root of the data.
        '''
        super(VQAFeatureDataset, self).__init__()
        assert name in ['train', 'val']

        ans2label_path = os.path.join(dataroot, 'cache', 'trainval_ans2label.pkl')
        label2ans_path = os.path.join(dataroot, 'cache', 'trainval_label2ans.pkl')
        self.ans2label = pickle.load(open(ans2label_path, 'rb'))
        self.label2ans = pickle.load(open(label2ans_path, 'rb'))
        self.num_ans_candidates = len(self.ans2label)

        self.dictionary = dictionary

        self.img_id2idx = pickle.load(
            open(os.path.join(dataroot, '{}36_imgid2idx.pkl'.format(name)), 'rb')
        )

        print('Loading features from h5 file')
        h5_path = os.path.join(dataroot, '{}36.hdf5'.format(name))
        with h5py.File(h5_path, 'r') as hf:
            self.features = np.array(hf.get('image_features'))
            self.spatials = np.array(hf.get('spatial_features'))

        self.entries = _load_dataset(dataroot, name, self.img_id2idx)
        self.tokenize()
        self.tensorize()

        self.v_dim = self.features.size(2)
        self.s_dim = self.spatials.size(2)

    def tokenize(self, max_length=14):
        '''
        Tokenizes the questions.

        This will add q_token in each entry of the dataset.
        -1 represent nil, and should be treated as padding_idx in embdding.
        '''
        for entry in self.entries:
            tokens = self.dictionary.tokenize(entry['question'], False)
            tokens = tokens[:max_length]
            if len(tokens) < max_length:
                padding = [self.dictionary.padding_idx] * (max_length - len(tokens))
                tokens = padding + tokens
            utils.assert_eq(len(tokens), max_length)
            entry['q_token'] = tokens

    def tensorize(self):
        self.features = torch.from_numpy(self.features)
        self.spatials = torch.from_numpy(self.spatials)

        for entry in self.entries:
            question = torch.from_numpy(np.array(entry['q_token']))
            entry['q_token'] = question

            answer = entry['answer']
            labels = np.array(answer['labels'])
            scores = np.array(answer['scores'], dtype=np.float32)

            if len(labels):
                labels = torch.from_numpy(labels)
                scores = torch.from_numpy(scores)
                entry['answer']['labels'] = labels
                entry['answer']['scores'] = scores
            else:
                entry['answer']['labels'] = None
                entry['answer']['scores'] = None

    def __getitem__(self, item):
        entry = self.entries[item]
        features = self.features[entry['image']]
        spatials = self.spatials[entry['image']]
        question = entry['q_token']
        answer = entry['answer']
        labels = answer['labels']
        scores = answer['scores']
        target = torch.zeros(self.num_ans_candidates) #type: torch.Tensor
        if labels is not None:
            target.scatter_(0, labels, scores)

        return features, spatials, question, target

    def __len__(self):
       return len(self.entries)










